# -*- coding: UTF-8 -*-

import numpy as np
import random
import math

import cv2

# 全局点索引
PtId = []


def pointToLineDist(point, line):
    """
    此处的line = [a, b, c]
    直线方程为: ax + by + c = 0
    :param point 直线外的一点 [x, y]
    :param line  直线方程参数数组 [a, b, c]
    :return float 类型的数据
    """
    # 不满足直线方程
    if line[0] < 0.000001 and line[1] < 0.000001:
        return 0.0

    # 算式分子
    numerator = float(math.fabs(line[0]*point[0] + line[1]*point[1] + line[2]))
    # 算式分母
    denominator = float(math.sqrt((line[0]*line[0])+(line[1]*line[1])))

    return float(numerator / denominator)


def makeLine(point1, point2):
    """
    由两个点组成直线方程ax + by + c = 0(a >= 0)
    :param point1 点1 [x, y]
    :param point2 点2 [x, y]
    :return 直线方程参数数组
    """
    line = [0., 0., 0.]
    sign = 1.0
    x1, y1 = point1[0], point1[1]
    x2, y2 = point2[0], point2[1]

    line[0] = y2 - y1
    if line[0] < 0:
        sign = (-1.0)
        line[0] = float(sign * line[0])

    line[1] = float(sign * (x1 - x2))
    line[2] = float(sign * (y1 * x2 - x1 * y2))

    return line


def fitLineRansac(points, iterations=1000, sigma=1.0, k_min=-7, k_max=7):
    """
    RANSAC 拟合2D 直线
    :param points:输入点集,numpy [points_num,1,2],np.float32
    :param iterations:迭代次数
    :param sigma:数据和模型之间可接受的差值,车道线像素宽带一般为10左右
                （Parameter use to compute the fitting score）
    :param k_min:
    :param k_max:k_min/k_max--拟合的直线斜率的取值范围.
                考虑到左右车道线在图像中的斜率位于一定范围内，
                添加此参数，同时可以避免检测垂线和水平线
    :return:拟合的直线参数,It is a vector of 4 floats
                (vx, vy, x0, y0) where (vx, vy) is a normalized
                vector collinear to the line and (x0, y0) is some
                point on the line.
    """
    line = [0, 0, 0, 0]
    points_num = points.shape[0]

    # 不足两个点，无法拟合直线
    if points_num < 2:
        return line

    #
    bestScore = -1
    for k in range(iterations):
        i1, i2 = random.sample(range(points_num), 2)
        p1 = points[i1][0]
        p2 = points[i2][0]

        # 两个点坐标相同时不能满足直线要求
        if p1[0] == p2[0] and p1[1] == p2[1]:
            continue

        dp = p1 - p2  # 直线的方向向量
        dp *= 1./np.linalg.norm(dp)  # 除以模长，进行归一化

        score = 0

        # 判断是垂直X轴的直线
        if math.fabs(dp[0]) < 0.00001:
            continue

        a = dp[1] / dp[0]

        if (a <= k_max) and (a >= k_min):
            for i in range(points_num):
                v = points[i][0] - p1
                dis = v[1]*dp[0] - v[0]*dp[1]  # 向量a与b叉乘/向量b的摸.||b||=1./norm(dp)
                # score += math.exp(-0.5*dis*dis/(sigma*sigma))误差定义方式的一种
                if math.fabs(dis) < sigma:
                    score += 1
                    PtId[i] = 1

        if score > bestScore:
            line = [dp[0], dp[1], p1[0], p1[1]]
            bestScore = score

    return line


def main():
    """对ransac拟合直线方法的测试"""
    image = np.ones([720, 1280, 3], dtype=np.ubyte) * 125

    # 以车道线参数为(0.7657, -0.6432, 534, 548)生成一系列点
    k = -0.6432 / 0.7657
    b1 = 548 - k * 534
    b2 = 548 - k * 635

    # 初始化空的点数组
    points = []

    # 第一条车道线
    for i in range(360, 720, 10):
        point = (int((i - b1) / k), i)
        points.append(point)

    # 第二条车道线
    for i in range(360, 720, 10):
        point = (int((i - b2) / k), i)
        points.append(point)

    # 在直线旁边加入机噪声
    for i in range(360, 720, 10):
        x = int((i - b1) / k)
        x = random.sample(range(x - 10, x + 10), 1)
        y = i
        y = random.sample(range(y - 30, y + 30), 1)

        point = (x[0], y[0])
        points.append(point)

    # # 加入噪声
    # for i in range(0, 720, 20):
    #     x = random.sample(range(1, 640), 1)
    #     y = random.sample(range(1, 360), 1)
    #     point = (x[0], y[0])
    #     points.append(point)

    # 在图像中画出点
    for point in points:
        # 画实心圆
        cv2.circle(image, point, 2, (0, 0, 0), -1)

    # 点数组转换成numpy格式
    points = np.array(points).astype(np.float32)
    points = points[:, np.newaxis, :]

    tmpPoints = []

    PtNum = points.shape[0]

    for i in range(PtNum):
        PtId.append(0)

    # RANSAC 拟合
    if 1:
        [vx, vy, x, y] = fitLineRansac(points, 1000, 5)
        k = float(vy) / float(vx)  # 直线斜率
        b = -k * x + y

        p1_y = 720
        p1_x = (p1_y - b) / k
        p2_y = 360
        p2_x = (p2_y - b) / k

        p1 = (int(p1_x), int(p1_y))
        p2 = (int(p2_x), int(p2_y))

        # cv2.line(image, p1, p2, (0, 255, 0), 2)

        line1 = [float(vy), float(-1.0*vx), float(vx*b)]
        if line1[0] < 0:
            line1[0] = line1[0] * (-1.0)
            line1[1] = line1[1] * (-1.0)
            line1[2] = line1[2] * (-1.0)

        # 画出RANSAC选出的内点
        for i in range(len(PtId)):
            # print(i, '====', index[i])
            if PtId[i] == 1:
                Pt2LineDist = pointToLineDist(points[i][0], line1)
                # print(i, " -----> ", Pt2LineDist)
                pt = (int(points[i][0][0]), int(points[i][0][1]))

                if Pt2LineDist < 15.0:
                    # tmpPoints = np.delete(tmpPoints, i, axis=1)
                    cv2.circle(image, pt, 3, (0, 0, 255), 2)
                # 将不满足条件的点保存
                else:
                    tmpPoints.append(pt)

    tmpPoints = np.array(tmpPoints).astype(np.float32)
    tmpPoints = tmpPoints[:, np.newaxis, :]
    # 对剩下的点集进行 RANSAC 拟合
    if 1:
        [vx, vy, x, y] = fitLineRansac(tmpPoints, 1000, 5)
        k = float(vy) / float(vx)  # 直线斜率
        b = -k * x + y

        p1_y = 720
        p1_x = (p1_y - b) / k
        p2_y = 360
        p2_x = (p2_y - b) / k

        p1 = (int(p1_x), int(p1_y))
        p2 = (int(p2_x), int(p2_y))

        cv2.line(image, p1, p2, (0, 255, 0), 2)

    # 最小二乘法拟合
    if 1:
        # [vx, vy, x, y] = cv2.fitLine(points, cv2.DIST_L2, 0, 0.1, 0.01)
        [vx, vy, x, y] = cv2.fitLine(tmpPoints, cv2.DIST_WELSCH, 0, 0.1, 0.01)
        k = float(vy) / float(vx)  # 直线斜率
        b = -k * x + y

        p1_y = 720
        p1_x = (p1_y - b) / k
        p2_y = 360
        p2_x = (p2_y - b) / k

        p1 = (int(p1_x), int(p1_y))
        p2 = (int(p2_x), int(p2_y))

        cv2.line(image, p1, p2, (0, 0, 255), 2)

    cv2.imshow('image', image)
    cv2.waitKey(0)
    cv2.destroyWindow('image')


# 运行此脚本直接调用测试函数
if __name__ == '__main__':
    main()

